package main

import (
	"bufio"
	"context"
	"crypto/md5"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"
)

// MCP协议常量
const (
	MCPVersion      = "2024-11-05"
	ServerName      = "drcom-server"
	ServerVersion   = "1.0.0"
	ProtocolVersion = "2024-11-05"
)

// 客户端验证相关结构
type ClientAuthConfig struct {
	Enabled            bool     `json:"enabled"`
	RequireAPIKey      bool     `json:"require_api_key"`
	RequireSignature   bool     `json:"require_signature"`
	RequireIPWhitelist bool     `json:"require_ip_whitelist"`
	APIKeys            []string `json:"api_keys"`
	SignatureSecret    string   `json:"signature_secret"`
	IPWhitelist        []string `json:"ip_whitelist"`
	MaxRetries         int      `json:"max_retries"`
	BanDuration        int      `json:"ban_duration_minutes"`
}

type AuthenticationInfo struct {
	APIKey    string `json:"api_key,omitempty"`
	Signature string `json:"signature,omitempty"`
	Timestamp string `json:"timestamp,omitempty"`
	ClientIP  string `json:"client_ip,omitempty"`
	UserAgent string `json:"user_agent,omitempty"`
}

type ClientSession struct {
	ID              string    `json:"id"`
	ClientIP        string    `json:"client_ip"`
	AuthTime        time.Time `json:"auth_time"`
	LastActivity    time.Time `json:"last_activity"`
	RequestCount    int       `json:"request_count"`
	IsAuthenticated bool      `json:"is_authenticated"`
	AuthMethod      string    `json:"auth_method"`
	FailedAttempts  int       `json:"failed_attempts"`
	IsBanned        bool      `json:"is_banned"`
	BanUntil        time.Time `json:"ban_until,omitempty"`
}

// 验证管理器
type AuthManager struct {
	config   ClientAuthConfig
	sessions map[string]*ClientSession
	mutex    sync.RWMutex
}

// MCP消息结构
type MCPRequest struct {
	JsonRPC string      `json:"jsonrpc"`
	ID      interface{} `json:"id,omitempty"`
	Method  string      `json:"method"`
	Params  interface{} `json:"params,omitempty"`
}

type MCPResponse struct {
	JsonRPC string      `json:"jsonrpc"`
	ID      interface{} `json:"id,omitempty"`
	Result  interface{} `json:"result,omitempty"`
	Error   *MCPError   `json:"error,omitempty"`
}

type MCPError struct {
	Code    int         `json:"code"`
	Message string      `json:"message"`
	Data    interface{} `json:"data,omitempty"`
}

// 初始化相关结构
type InitializeParams struct {
	ProtocolVersion string              `json:"protocolVersion"`
	Capabilities    ClientCapabilities  `json:"capabilities"`
	ClientInfo      ClientInfo          `json:"clientInfo"`
	Authentication  *AuthenticationInfo `json:"authentication,omitempty"`
}

type ClientCapabilities struct {
	Roots    *RootsCapability    `json:"roots,omitempty"`
	Sampling *SamplingCapability `json:"sampling,omitempty"`
}

type RootsCapability struct {
	ListChanged bool `json:"listChanged,omitempty"`
}

type SamplingCapability struct{}

type ClientInfo struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

type ServerCapabilities struct {
	Logging   *LoggingCapability   `json:"logging,omitempty"`
	Tools     *ToolsCapability     `json:"tools,omitempty"`
	Resources *ResourcesCapability `json:"resources,omitempty"`
	Prompts   *PromptsCapability   `json:"prompts,omitempty"`
}

type LoggingCapability struct{}
type ToolsCapability struct {
	ListChanged bool `json:"listChanged,omitempty"`
}
type ResourcesCapability struct {
	Subscribe   bool `json:"subscribe,omitempty"`
	ListChanged bool `json:"listChanged,omitempty"`
}
type PromptsCapability struct {
	ListChanged bool `json:"listChanged,omitempty"`
}

type ServerInfo struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

type InitializeResult struct {
	ProtocolVersion string             `json:"protocolVersion"`
	Capabilities    ServerCapabilities `json:"capabilities"`
	ServerInfo      ServerInfo         `json:"serverInfo"`
}

// 工具相关结构
type Tool struct {
	Name        string      `json:"name"`
	Description string      `json:"description"`
	InputSchema interface{} `json:"inputSchema"`
}

type ToolsListResult struct {
	Tools []Tool `json:"tools"`
}

type CallToolParams struct {
	Name      string                 `json:"name"`
	Arguments map[string]interface{} `json:"arguments,omitempty"`
}

type ToolResult struct {
	Content []ContentItem `json:"content"`
	IsError bool          `json:"isError,omitempty"`
}

type ContentItem struct {
	Type string `json:"type"`
	Text string `json:"text"`
}

// 业务相关结构（根据Drcom API文档配置）
// API调用方式：
// URL: http://demo.doctorcom.com:8088/DrcomSrv/DrcomService
// 方法: GET
// 参数: iusername, business(BASE64编码的JSON指令), timestamp, sign(MD5签名)
// 签名规则: MD5(business + timestamp + signKey)
type ActionConfig struct {
	Method          string            `json:"method"`
	URL             string            `json:"url"`
	Template        string            `json:"template"`
	Params          map[string]string `json:"params"`
	InstructionCode string            `json:"instruction_code"`
	Iusername       string            `json:"iusername"`
	SignKey         string            `json:"sign_key"`
}

var mcpConfig = map[string]ActionConfig{
	"query_user_usage": {
		Method:          "GET",
		URL:             "http://demo.doctorcom.com:8088/DrcomSrv/DrcomService",
		Template:        "用户{{.account}}使用情况：\n• 用户状态：{{.user_state_desc}}\n• 状态时间：{{.state_time}}\n• 储值余额：{{.balance}}元\n• 本期已使用费用：{{.use_money}}元\n• 本期已使用时长：{{.use_time_hours}}小时\n• 本期已使用流量：{{.use_flow_gb}}GB\n• 起始计费日：{{.start_date}}\n• 状态备注：{{.state_memo}}",
		InstructionCode: "091",
		Iusername:       "ykt",
		SignKey:         "RORB1JKy4pSSAJtm",
		Params: map[string]string{
			"account": "user_account",
		},
	},
}

type BusinessResponse map[string]interface{}

// 添加并发处理相关的结构
type RequestJob struct {
	Request *MCPRequest
	Done    chan bool
}

type WorkerPool struct {
	WorkerCount int
	JobQueue    chan RequestJob
	Workers     []*Worker
	Server      *MCPServer
	ctx         context.Context
	cancel      context.CancelFunc
	wg          sync.WaitGroup
}

type Worker struct {
	ID       int
	JobQueue chan RequestJob
	Quit     chan bool
	Server   *MCPServer
	wg       *sync.WaitGroup
}

// 添加HTTP客户端池
type HTTPClientPool struct {
	clients chan *http.Client
	maxSize int
	timeout time.Duration
}

// 批量查询请求结构
type BatchUserQuery struct {
	UserAccounts []string `json:"user_accounts"`
}

type BatchQueryResult struct {
	Results map[string]interface{} `json:"results"`
	Errors  map[string]string      `json:"errors"`
}

func main() {
	// 设置日志输出到stderr
	log.SetOutput(os.Stderr)

	// 显示启动信息
	log.Println("=== Drcom MCP Server 启动 ===")
	log.Printf("服务器名称: %s", ServerName)
	log.Printf("服务器版本: %s", ServerVersion)
	log.Printf("协议版本: %s", ProtocolVersion)
	log.Println("配置信息:")

	for name, config := range mcpConfig {
		log.Printf("  [%s] URL: %s", name, config.URL)
		log.Printf("  [%s] 指令代码: %s", name, config.InstructionCode)
		log.Printf("  [%s] 接口用户: %s", name, config.Iusername)
	}

	log.Println("支持的工具:")
	log.Println("  - natural_language_query: 自然语言查询")
	log.Println("  - query_user_usage: 查询用户使用量")
	log.Println("  - batch_query_users: 批量查询用户")

	// 预加载验证配置以显示状态
	authConfig := loadAuthConfig()
	log.Println("客户端验证配置:")
	if authConfig.Enabled {
		log.Printf("  状态: 已启用")
		log.Printf("  API Key验证: %t", authConfig.RequireAPIKey)
		log.Printf("  签名验证: %t", authConfig.RequireSignature)
		log.Printf("  IP白名单验证: %t", authConfig.RequireIPWhitelist)
		log.Printf("  最大重试次数: %d", authConfig.MaxRetries)
		log.Printf("  封禁时长: %d分钟", authConfig.BanDuration)
		if authConfig.RequireAPIKey {
			log.Printf("  已配置API密钥数量: %d", len(authConfig.APIKeys))
		}
		if authConfig.RequireIPWhitelist {
			log.Printf("  IP白名单数量: %d", len(authConfig.IPWhitelist))
		}
	} else {
		log.Printf("  状态: 已禁用")
	}

	log.Println()
	log.Println("MCP服务器已准备就绪，等待客户端连接...")
	log.Println("=====================================")

	server := NewMCPServer()
	server.Run()
}

type MCPServer struct {
	scanner     *bufio.Scanner
	workerPool  *WorkerPool
	httpPool    *HTTPClientPool
	mu          sync.RWMutex // 保护并发访问
	authManager *AuthManager // 验证管理器
}

func NewMCPServer() *MCPServer {
	// 创建HTTP客户端池
	httpPool := &HTTPClientPool{
		clients: make(chan *http.Client, 10), // 最多10个HTTP客户端
		maxSize: 10,
		timeout: 30 * time.Second,
	}

	// 初始化HTTP客户端池
	for i := 0; i < 10; i++ {
		client := &http.Client{
			Timeout: httpPool.timeout,
			Transport: &http.Transport{
				MaxIdleConns:        100,
				MaxIdleConnsPerHost: 10,
				IdleConnTimeout:     90 * time.Second,
			},
		}
		httpPool.clients <- client
	}

	// 创建验证管理器
	authConfig := loadAuthConfig()
	authManager := &AuthManager{
		config:   authConfig,
		sessions: make(map[string]*ClientSession),
	}

	server := &MCPServer{
		scanner:     bufio.NewScanner(os.Stdin),
		httpPool:    httpPool,
		authManager: authManager,
	}

	// 创建工作池
	ctx, cancel := context.WithCancel(context.Background())
	workerPool := &WorkerPool{
		WorkerCount: 5,                          // 5个工作协程
		JobQueue:    make(chan RequestJob, 100), // 缓冲100个请求
		Server:      server,
		ctx:         ctx,
		cancel:      cancel,
	}

	server.workerPool = workerPool
	return server
}

// 获取HTTP客户端
func (pool *HTTPClientPool) GetClient() *http.Client {
	select {
	case client := <-pool.clients:
		return client
	default:
		// 如果池中没有可用客户端，创建新的
		return &http.Client{
			Timeout: pool.timeout,
			Transport: &http.Transport{
				MaxIdleConns:        100,
				MaxIdleConnsPerHost: 10,
				IdleConnTimeout:     90 * time.Second,
			},
		}
	}
}

// 归还HTTP客户端
func (pool *HTTPClientPool) PutClient(client *http.Client) {
	select {
	case pool.clients <- client:
	default:
		// 如果池已满，则丢弃客户端
	}
}

// 创建工作者
func NewWorker(id int, jobQueue chan RequestJob, server *MCPServer, wg *sync.WaitGroup) *Worker {
	return &Worker{
		ID:       id,
		JobQueue: jobQueue,
		Quit:     make(chan bool),
		Server:   server,
		wg:       wg,
	}
}

// 工作者处理逻辑
func (w *Worker) Start() {
	w.wg.Add(1)
	go func() {
		defer w.wg.Done()
		log.Printf("Worker %d 启动", w.ID)

		for {
			select {
			case job := <-w.JobQueue:
				log.Printf("Worker %d 处理请求: %s", w.ID, job.Request.Method)
				w.Server.handleRequestConcurrent(job.Request)
				job.Done <- true

			case <-w.Quit:
				log.Printf("Worker %d 停止", w.ID)
				return
			}
		}
	}()
}

// 停止工作者
func (w *Worker) Stop() {
	w.Quit <- true
}

// 启动工作池
func (wp *WorkerPool) Start() {
	wp.Workers = make([]*Worker, wp.WorkerCount)

	for i := 0; i < wp.WorkerCount; i++ {
		worker := NewWorker(i+1, wp.JobQueue, wp.Server, &wp.wg)
		wp.Workers[i] = worker
		worker.Start()
	}

	log.Printf("工作池启动，共 %d 个工作协程", wp.WorkerCount)
}

// 停止工作池
func (wp *WorkerPool) Stop() {
	log.Println("正在停止工作池...")
	wp.cancel()

	// 停止所有工作者
	for _, worker := range wp.Workers {
		worker.Stop()
	}

	// 等待所有工作者完成
	wp.wg.Wait()
	log.Println("工作池已停止")
}

// 提交任务到工作池
func (wp *WorkerPool) Submit(request *MCPRequest) {
	job := RequestJob{
		Request: request,
		Done:    make(chan bool, 1),
	}

	select {
	case wp.JobQueue <- job:
		// 等待任务完成（可选，根据需要调整）
		<-job.Done
	case <-wp.ctx.Done():
		log.Println("工作池已关闭，拒绝新任务")
	default:
		log.Println("工作池队列已满，直接处理请求")
		wp.Server.handleRequestConcurrent(request)
	}
}

func (s *MCPServer) Run() {
	log.SetOutput(os.Stderr) // 将日志输出到stderr，避免干扰MCP通信

	// 启动工作池
	s.workerPool.Start()
	defer s.workerPool.Stop()

	for s.scanner.Scan() {
		line := s.scanner.Text()
		if line == "" {
			continue
		}

		var request MCPRequest
		if err := json.Unmarshal([]byte(line), &request); err != nil {
			s.sendError(nil, -32700, "Parse error", err.Error())
			continue
		}

		// 对于某些请求（如初始化），仍然同步处理
		// 对于工具调用等，使用并发处理
		if request.Method == "initialize" || request.Method == "notifications/initialized" {
			s.handleRequest(&request)
		} else {
			// 提交到工作池进行并发处理
			go s.workerPool.Submit(&request)
		}
	}

	if err := s.scanner.Err(); err != nil {
		log.Printf("Scanner error: %v", err)
	}
}

// 并发处理请求的方法
func (s *MCPServer) handleRequestConcurrent(request *MCPRequest) {
	log.Printf("并发处理MCP请求: %s (ID: %v)", request.Method, request.ID)

	switch request.Method {
	case "tools/list":
		log.Printf("处理工具列表请求")
		s.handleToolsList(request)
	case "tools/call":
		log.Printf("处理工具调用请求")
		s.handleToolsCallConcurrent(request)
	default:
		log.Printf("未知请求方法: %s", request.Method)
		s.sendError(request.ID, -32601, "Method not found", fmt.Sprintf("Unknown method: %s", request.Method))
	}
}

// 并发处理工具调用
func (s *MCPServer) handleToolsCallConcurrent(request *MCPRequest) {
	paramsBytes, err := json.Marshal(request.Params)
	if err != nil {
		log.Printf("工具调用参数序列化失败: %v", err)
		s.sendError(request.ID, -32602, "Invalid params", err.Error())
		return
	}

	var params CallToolParams
	if err := json.Unmarshal(paramsBytes, &params); err != nil {
		log.Printf("工具调用参数解析失败: %v", err)
		s.sendError(request.ID, -32602, "Invalid params", err.Error())
		return
	}

	log.Printf("并发调用工具: %s, 参数: %+v", params.Name, params.Arguments)

	switch params.Name {
	case "natural_language_query":
		log.Printf("执行自然语言查询")
		s.handleNaturalLanguageQuery(request.ID, params.Arguments)
	case "query_user_usage":
		log.Printf("执行用户使用量查询")
		s.handleQueryUserUsageConcurrent(request.ID, params.Arguments)
	case "batch_query_users":
		log.Printf("执行批量用户查询")
		s.handleBatchQueryUsers(request.ID, params.Arguments)
	default:
		log.Printf("未知工具: %s", params.Name)
		s.sendError(request.ID, -32602, "Unknown tool", fmt.Sprintf("Tool '%s' not found", params.Name))
	}
}

// 并发处理用户使用量查询
func (s *MCPServer) handleQueryUserUsageConcurrent(requestID interface{}, args map[string]interface{}) {
	userAccount, ok := args["user_account"].(string)
	if !ok {
		s.sendError(requestID, -32602, "Invalid arguments", "user_account is required and must be a string")
		return
	}

	// 验证用户账号格式
	if userAccount == "" {
		s.sendError(requestID, -32602, "Invalid arguments", "user_account cannot be empty")
		return
	}

	config := mcpConfig["query_user_usage"]
	params := map[string]string{"user_account": userAccount}

	// 使用协程池中的HTTP客户端进行API调用
	bizResp, err := s.callDrcomAPIConcurrent(config, params)
	if err != nil {
		log.Printf("Business API error: %v", err)
		s.sendToolError(requestID, fmt.Sprintf("业务API调用失败: %v", err))
		return
	}

	// 检查响应结果
	resultCode, ok := bizResp["result"].(string)
	if !ok || resultCode != "E00" {
		resultMsg := "未知错误"
		if msg, exists := bizResp["result_msg"].(string); exists {
			resultMsg = msg
		}
		s.sendToolError(requestID, fmt.Sprintf("API调用失败: %s (代码: %s)", resultMsg, resultCode))
		return
	}

	// 提取用户数据
	list, ok := bizResp["list"].([]interface{})
	if !ok || len(list) == 0 {
		s.sendToolError(requestID, "未找到用户数据")
		return
	}

	userData, ok := list[0].(map[string]interface{})
	if !ok {
		s.sendToolError(requestID, "用户数据格式错误")
		return
	}

	// 添加用户账号到数据中
	userData["account"] = userAccount

	// 生成响应
	response := generateResponse(config.Template, userData)

	result := ToolResult{
		Content: []ContentItem{
			{
				Type: "text",
				Text: response,
			},
		},
		IsError: false,
	}

	s.sendResponse(requestID, result)
}

// 新增：批量查询用户信息
func (s *MCPServer) handleBatchQueryUsers(requestID interface{}, args map[string]interface{}) {
	userAccountsInterface, ok := args["user_accounts"]
	if !ok {
		s.sendError(requestID, -32602, "Invalid arguments", "user_accounts is required")
		return
	}

	// 解析用户账号列表
	var userAccounts []string
	switch v := userAccountsInterface.(type) {
	case []interface{}:
		for _, account := range v {
			if accountStr, ok := account.(string); ok {
				userAccounts = append(userAccounts, accountStr)
			}
		}
	case []string:
		userAccounts = v
	default:
		s.sendError(requestID, -32602, "Invalid arguments", "user_accounts must be an array of strings")
		return
	}

	if len(userAccounts) == 0 {
		s.sendError(requestID, -32602, "Invalid arguments", "user_accounts cannot be empty")
		return
	}

	log.Printf("批量查询用户: %v", userAccounts)

	// 并发查询所有用户
	results := make(map[string]interface{})
	errors := make(map[string]string)
	var wg sync.WaitGroup
	var mu sync.Mutex

	config := mcpConfig["query_user_usage"]

	for _, userAccount := range userAccounts {
		wg.Add(1)
		go func(account string) {
			defer wg.Done()

			params := map[string]string{"user_account": account}
			bizResp, err := s.callDrcomAPIConcurrent(config, params)

			mu.Lock()
			defer mu.Unlock()

			if err != nil {
				errors[account] = fmt.Sprintf("API调用失败: %v", err)
				return
			}

			resultCode, ok := bizResp["result"].(string)
			if !ok || resultCode != "E00" {
				resultMsg := "未知错误"
				if msg, exists := bizResp["result_msg"].(string); exists {
					resultMsg = msg
				}
				errors[account] = fmt.Sprintf("API调用失败: %s (代码: %s)", resultMsg, resultCode)
				return
			}

			list, ok := bizResp["list"].([]interface{})
			if !ok || len(list) == 0 {
				errors[account] = "未找到用户数据"
				return
			}

			userData, ok := list[0].(map[string]interface{})
			if !ok {
				errors[account] = "用户数据格式错误"
				return
			}

			userData["account"] = account
			response := generateResponse(config.Template, userData)
			results[account] = response

		}(userAccount)
	}

	// 等待所有查询完成
	wg.Wait()

	// 生成响应文本
	var responseText strings.Builder
	responseText.WriteString(fmt.Sprintf("批量查询完成，共查询 %d 个用户：\n\n", len(userAccounts)))

	for account, result := range results {
		responseText.WriteString(fmt.Sprintf("=== %s ===\n", account))
		responseText.WriteString(fmt.Sprintf("%v\n\n", result))
	}

	if len(errors) > 0 {
		responseText.WriteString("查询错误：\n")
		for account, errMsg := range errors {
			responseText.WriteString(fmt.Sprintf("- %s: %s\n", account, errMsg))
		}
	}

	result := ToolResult{
		Content: []ContentItem{
			{
				Type: "text",
				Text: responseText.String(),
			},
		},
		IsError: len(errors) > 0 && len(results) == 0, // 只有在所有查询都失败时才标记为错误
	}

	s.sendResponse(requestID, result)
}

// 使用连接池的API调用方法
func (s *MCPServer) callDrcomAPIConcurrent(config ActionConfig, params map[string]string) (BusinessResponse, error) {
	log.Printf("\n=== 开始并发API调用 [指令%s] ===", config.InstructionCode)
	log.Printf("目标URL: %s", config.URL)
	log.Printf("接口用户: %s", config.Iusername)
	log.Printf("请求参数: %+v", params)

	// 构建指令JSON
	instruction := map[string]interface{}{
		"code": config.InstructionCode,
	}

	var instructionJSON []byte
	var err error

	switch config.InstructionCode {
	case "091":
		var userAccount string
		if account, ok := params["user_account"]; ok {
			userAccount = account
			log.Printf("用户账号: %s", userAccount)
		} else {
			return nil, fmt.Errorf("未能从查询中解析出用户账号，请提供用户名")
		}
		instructionJSONStr := fmt.Sprintf(`{"code":"%s","account":"%s"}`, config.InstructionCode, userAccount)
		instructionJSON = []byte(instructionJSONStr)
		log.Printf("指令JSON: %s", instructionJSONStr)
	default:
		for paramKey, valueKey := range config.Params {
			if val, ok := params[valueKey]; ok {
				instruction[paramKey] = val
				log.Printf("添加参数: %s = %s", paramKey, val)
			}
		}
		instructionJSON, err = json.Marshal(instruction)
		if err != nil {
			log.Printf("指令JSON序列化失败: %v", err)
			return nil, fmt.Errorf("failed to marshal instruction: %v", err)
		}
		log.Printf("指令JSON: %s", string(instructionJSON))
	}

	// 对指令进行BASE64编码
	business := base64.StdEncoding.EncodeToString(instructionJSON)
	log.Printf("Business参数(BASE64): %s", business)

	// 生成时间戳
	timestamp := time.Now().Format("20060102150405")
	log.Printf("时间戳: %s", timestamp)

	// 生成签名
	signData := business + timestamp + config.SignKey
	hasher := md5.New()
	hasher.Write([]byte(signData))
	sign := hex.EncodeToString(hasher.Sum(nil))
	log.Printf("签名数据: %s", signData)
	log.Printf("MD5签名: %s", sign)

	// 构建请求URL
	baseURL, err := url.Parse(config.URL)
	if err != nil {
		log.Printf("URL解析失败: %v", err)
		return nil, fmt.Errorf("invalid URL: %v", err)
	}

	query := baseURL.Query()
	query.Add("iusername", config.Iusername)
	query.Add("business", business)
	query.Add("timestamp", timestamp)
	query.Add("sign", sign)
	baseURL.RawQuery = query.Encode()

	log.Printf("完整请求URL: %s", baseURL.String())

	// 从连接池获取HTTP客户端
	client := s.httpPool.GetClient()
	defer s.httpPool.PutClient(client)

	// 创建HTTP请求
	req, err := http.NewRequest(config.Method, baseURL.String(), nil)
	if err != nil {
		log.Printf("创建HTTP请求失败: %v", err)
		return nil, err
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded;charset=UTF-8")
	log.Printf("请求方法: %s", config.Method)

	// 发送请求
	log.Printf("发送并发HTTP请求...")
	resp, err := client.Do(req)
	if err != nil {
		log.Printf("HTTP请求失败: %v", err)
		return nil, err
	}
	defer resp.Body.Close()

	log.Printf("响应状态码: %d", resp.StatusCode)

	if resp.StatusCode != http.StatusOK {
		body, _ := ioutil.ReadAll(resp.Body)
		log.Printf("API错误响应: %s", string(body))
		return nil, fmt.Errorf("API error: status=%d, body=%s", resp.StatusCode, body)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Printf("读取响应体失败: %v", err)
		return nil, err
	}

	log.Printf("原始响应体: %s", string(body))

	var result BusinessResponse
	if err := json.Unmarshal(body, &result); err != nil {
		log.Printf("JSON解析失败: %v", err)
		return nil, fmt.Errorf("failed to parse API response: %v, body: %s", err, string(body))
	}

	log.Printf("=== 并发API调用完成 ===\n")
	return result, nil
}

func (s *MCPServer) handleRequest(request *MCPRequest) {
	log.Printf("收到MCP请求: %s (ID: %v)", request.Method, request.ID)

	switch request.Method {
	case "initialize":
		log.Printf("处理初始化请求")
		s.handleInitialize(request)
	case "notifications/initialized":
		log.Printf("收到客户端初始化完成通知")
		// 客户端初始化完成通知，无需响应
		return
	case "tools/list":
		log.Printf("处理工具列表请求")
		s.handleToolsList(request)
	case "tools/call":
		log.Printf("处理工具调用请求")
		s.handleToolsCall(request)
	default:
		log.Printf("未知请求方法: %s", request.Method)
		s.sendError(request.ID, -32601, "Method not found", fmt.Sprintf("Unknown method: %s", request.Method))
	}
}

func (s *MCPServer) handleInitialize(request *MCPRequest) {
	// 解析初始化参数
	paramsBytes, err := json.Marshal(request.Params)
	if err != nil {
		log.Printf("初始化参数序列化失败: %v", err)
		s.sendError(request.ID, -32602, "Invalid params", err.Error())
		return
	}

	var params InitializeParams
	if err := json.Unmarshal(paramsBytes, &params); err != nil {
		log.Printf("初始化参数解析失败: %v", err)
		s.sendError(request.ID, -32602, "Invalid params", err.Error())
		return
	}

	// 获取客户端IP
	clientIP := getClientIP()
	log.Printf("客户端初始化请求 - IP: %s, 客户端: %s %s",
		clientIP, params.ClientInfo.Name, params.ClientInfo.Version)

	// 执行客户端验证
	session, err := s.authManager.authenticateClient(params.Authentication, clientIP)
	if err != nil {
		log.Printf("客户端验证失败: %v", err)
		s.sendError(request.ID, -32001, "Authentication failed", err.Error())
		return
	}

	log.Printf("客户端验证成功 - 会话ID: %s, 方法: %s", session.ID, session.AuthMethod)

	// 启动会话清理协程（仅首次验证时启动）
	go func() {
		ticker := time.NewTicker(1 * time.Hour)
		defer ticker.Stop()
		for range ticker.C {
			s.authManager.cleanupExpiredSessions()
		}
	}()

	// 返回初始化结果
	result := InitializeResult{
		ProtocolVersion: ProtocolVersion,
		Capabilities: ServerCapabilities{
			Tools: &ToolsCapability{
				ListChanged: false,
			},
			Logging: &LoggingCapability{},
		},
		ServerInfo: ServerInfo{
			Name:    ServerName,
			Version: ServerVersion,
		},
	}

	s.sendResponse(request.ID, result)
}

func (s *MCPServer) handleToolsList(request *MCPRequest) {
	tools := []Tool{
		{
			Name:        "natural_language_query",
			Description: "使用自然语言查询用户信息，支持中文查询如'用户drset使用量'等",
			InputSchema: map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"query": map[string]interface{}{
						"type":        "string",
						"description": "自然语言查询，如'用户drset使用量'",
					},
				},
				"required": []string{"query"},
			},
		},
		{
			Name:        "query_user_usage",
			Description: "查询单个用户的使用量信息，包括费用、时长、流量和余额",
			InputSchema: map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"user_account": map[string]interface{}{
						"type":        "string",
						"description": "用户账号名称",
					},
				},
				"required": []string{"user_account"},
			},
		},
		{
			Name:        "batch_query_users",
			Description: "并发批量查询多个用户的使用量信息，提高查询效率",
			InputSchema: map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"user_accounts": map[string]interface{}{
						"type":        "array",
						"description": "用户账号名称列表",
						"items": map[string]interface{}{
							"type": "string",
						},
					},
				},
				"required": []string{"user_accounts"},
			},
		},
	}

	result := ToolsListResult{Tools: tools}
	s.sendResponse(request.ID, result)
}

func (s *MCPServer) handleToolsCall(request *MCPRequest) {
	paramsBytes, err := json.Marshal(request.Params)
	if err != nil {
		log.Printf("工具调用参数序列化失败: %v", err)
		s.sendError(request.ID, -32602, "Invalid params", err.Error())
		return
	}

	var params CallToolParams
	if err := json.Unmarshal(paramsBytes, &params); err != nil {
		log.Printf("工具调用参数解析失败: %v", err)
		s.sendError(request.ID, -32602, "Invalid params", err.Error())
		return
	}

	log.Printf("调用工具: %s, 参数: %+v", params.Name, params.Arguments)

	switch params.Name {
	case "natural_language_query":
		log.Printf("执行自然语言查询")
		s.handleNaturalLanguageQuery(request.ID, params.Arguments)
	case "query_user_usage":
		log.Printf("执行用户使用量查询")
		s.handleQueryUserUsage(request.ID, params.Arguments)
	case "batch_query_users":
		log.Printf("执行批量用户查询")
		s.handleBatchQueryUsers(request.ID, params.Arguments)
	default:
		log.Printf("未知工具: %s", params.Name)
		s.sendError(request.ID, -32602, "Unknown tool", fmt.Sprintf("Tool '%s' not found", params.Name))
	}
}

func (s *MCPServer) handleNaturalLanguageQuery(requestID interface{}, args map[string]interface{}) {
	query, ok := args["query"].(string)
	if !ok {
		s.sendError(requestID, -32602, "Invalid arguments", "query is required and must be a string")
		return
	}

	// 解析自然语言查询
	action, params := parseNaturalLanguage(query)
	if action == "" {
		s.sendToolError(requestID, "无法理解查询内容，请尝试使用如下格式：\n- '用户drset使用量'\n- '用户test001使用情况'\n- 或直接输入用户账号名称")
		return
	}

	// 根据解析结果调用相应的处理方法
	switch action {
	case "query_user_usage":
		s.handleQueryUserUsage(requestID, map[string]interface{}{"user_account": params["user_account"]})
	default:
		s.sendToolError(requestID, fmt.Sprintf("不支持的操作类型: %s", action))
	}
}

func (s *MCPServer) handleQueryUserUsage(requestID interface{}, args map[string]interface{}) {
	userAccount, ok := args["user_account"].(string)
	if !ok {
		s.sendError(requestID, -32602, "Invalid arguments", "user_account is required and must be a string")
		return
	}

	// 验证用户账号格式（移除数字限制，允许任意字符串）
	if userAccount == "" {
		s.sendError(requestID, -32602, "Invalid arguments", "user_account cannot be empty")
		return
	}

	config := mcpConfig["query_user_usage"]
	params := map[string]string{"user_account": userAccount}

	bizResp, err := callDrcomAPI(config, params)
	if err != nil {
		log.Printf("Business API error: %v", err)
		s.sendToolError(requestID, fmt.Sprintf("业务API调用失败: %v", err))
		return
	}

	// 检查响应结果
	resultCode, ok := bizResp["result"].(string)
	if !ok || resultCode != "E00" {
		resultMsg := "未知错误"
		if msg, exists := bizResp["result_msg"].(string); exists {
			resultMsg = msg
		}
		s.sendToolError(requestID, fmt.Sprintf("API调用失败: %s (代码: %s)", resultMsg, resultCode))
		return
	}

	// 提取用户数据
	list, ok := bizResp["list"].([]interface{})
	if !ok || len(list) == 0 {
		s.sendToolError(requestID, "未找到用户数据")
		return
	}

	userData, ok := list[0].(map[string]interface{})
	if !ok {
		s.sendToolError(requestID, "用户数据格式错误")
		return
	}

	// 添加用户账号到数据中
	userData["account"] = userAccount

	result := generateResponse(config.Template, userData)

	toolResult := ToolResult{
		Content: []ContentItem{
			{
				Type: "text",
				Text: result,
			},
		},
		IsError: false,
	}

	s.sendResponse(requestID, toolResult)
}

func (s *MCPServer) sendResponse(id interface{}, result interface{}) {
	response := MCPResponse{
		JsonRPC: "2.0",
		ID:      id,
		Result:  result,
	}

	s.sendMessage(response)
}

func (s *MCPServer) sendError(id interface{}, code int, message string, data interface{}) {
	response := MCPResponse{
		JsonRPC: "2.0",
		ID:      id,
		Error: &MCPError{
			Code:    code,
			Message: message,
			Data:    data,
		},
	}

	s.sendMessage(response)
}

func (s *MCPServer) sendToolError(id interface{}, errorMsg string) {
	toolResult := ToolResult{
		Content: []ContentItem{
			{
				Type: "text",
				Text: errorMsg,
			},
		},
		IsError: true,
	}

	s.sendResponse(id, toolResult)
}

func (s *MCPServer) sendMessage(message interface{}) {
	data, err := json.Marshal(message)
	if err != nil {
		log.Printf("Failed to marshal message: %v", err)
		return
	}

	fmt.Println(string(data))
}

// 以下是复用的业务逻辑代码

// 自然语言解析
func parseNaturalLanguage(query string) (string, map[string]string) {
	params := make(map[string]string)

	// 用户使用量查询
	if (strings.Contains(query, "使用量") || strings.Contains(query, "使用情况") ||
		strings.Contains(query, "已使用") || strings.Contains(query, "费用") ||
		strings.Contains(query, "时长") || strings.Contains(query, "流量") ||
		strings.Contains(query, "余额")) && strings.Contains(query, "用户") {
		// 匹配用户后跟任意字符（用户名），支持更多字符类型
		re := regexp.MustCompile(`用户([a-zA-Z0-9_\-\.@]+)`)
		matches := re.FindStringSubmatch(query)
		if len(matches) > 1 {
			params["user_account"] = matches[1]
			return "query_user_usage", params
		}

		// 如果正则匹配失败，尝试更宽松的匹配（用户和其后的非空白字符）
		re2 := regexp.MustCompile(`用户\s*([^\s，。！？,.\!?]+)`)
		matches2 := re2.FindStringSubmatch(query)
		if len(matches2) > 1 {
			params["user_account"] = matches2[1]
			return "query_user_usage", params
		}
	}

	// 直接输入用户账号名称的查询
	query = strings.TrimSpace(query)
	if query != "" && !strings.Contains(query, " ") {
		params["user_account"] = query
		// 默认查询使用量
		return "query_user_usage", params
	}

	return "", params
}

func callDrcomAPI(config ActionConfig, params map[string]string) (BusinessResponse, error) {
	log.Printf("\n=== 开始API调用 [指令%s] ===", config.InstructionCode)
	log.Printf("目标URL: %s", config.URL)
	log.Printf("接口用户: %s", config.Iusername)
	log.Printf("请求参数: %+v", params)

	// 构建指令JSON - 使用标准的JSON格式
	instruction := map[string]interface{}{
		"code": config.InstructionCode,
	}

	// 根据不同的指令代码添加相应的参数
	var instructionJSON []byte
	var err error

	switch config.InstructionCode {
	case "091":
		// 查询用户使用量，按照文档格式: {"code":"091","account":"用户名"}
		// 手动构建JSON确保字段顺序正确
		var userAccount string
		if account, ok := params["user_account"]; ok {
			userAccount = account
			log.Printf("用户账号: %s", userAccount)
		} else {
			// 如果没有提供用户账号，返回错误
			return nil, fmt.Errorf("未能从查询中解析出用户账号，请提供用户名")
		}
		// 确保字段顺序：code在前，account在后
		instructionJSONStr := fmt.Sprintf(`{"code":"%s","account":"%s"}`, config.InstructionCode, userAccount)
		instructionJSON = []byte(instructionJSONStr)
		log.Printf("指令JSON: %s", instructionJSONStr)
	default:
		// 其他指令类型，使用原有逻辑
		for paramKey, valueKey := range config.Params {
			if val, ok := params[valueKey]; ok {
				instruction[paramKey] = val
				log.Printf("添加参数: %s = %s", paramKey, val)
			}
		}
		// 将指令转换为JSON字符串
		instructionJSON, err = json.Marshal(instruction)
		if err != nil {
			log.Printf("指令JSON序列化失败: %v", err)
			return nil, fmt.Errorf("failed to marshal instruction: %v", err)
		}
		log.Printf("指令JSON: %s", string(instructionJSON))
	}

	// 对指令进行BASE64编码
	business := base64.StdEncoding.EncodeToString(instructionJSON)
	log.Printf("Business参数(BASE64): %s", business)

	// 生成时间戳
	timestamp := time.Now().Format("20060102150405")
	log.Printf("时间戳: %s", timestamp)

	// 生成签名：按照文档要求，对business + timestamp + 签名key进行MD5签名
	signData := business + timestamp + config.SignKey
	hasher := md5.New()
	hasher.Write([]byte(signData))
	sign := hex.EncodeToString(hasher.Sum(nil))
	log.Printf("签名数据: %s", signData)
	log.Printf("MD5签名: %s", sign)

	// 构建请求URL
	baseURL, err := url.Parse(config.URL)
	if err != nil {
		log.Printf("URL解析失败: %v", err)
		return nil, fmt.Errorf("invalid URL: %v", err)
	}

	// 添加查询参数
	query := baseURL.Query()
	query.Add("iusername", config.Iusername)
	query.Add("business", business)
	query.Add("timestamp", timestamp)
	query.Add("sign", sign)
	baseURL.RawQuery = query.Encode()

	log.Printf("完整请求URL: %s", baseURL.String())

	// 创建HTTP请求
	req, err := http.NewRequest(config.Method, baseURL.String(), nil)
	if err != nil {
		log.Printf("创建HTTP请求失败: %v", err)
		return nil, err
	}

	// 设置请求头
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded;charset=UTF-8")
	log.Printf("请求方法: %s", config.Method)
	log.Printf("请求头: %+v", req.Header)

	// 发送请求
	client := &http.Client{
		Timeout: 30 * time.Second, // 设置超时时间
	}
	log.Printf("发送HTTP请求...")
	resp, err := client.Do(req)
	if err != nil {
		log.Printf("HTTP请求失败: %v", err)
		return nil, err
	}
	defer resp.Body.Close()

	log.Printf("响应状态码: %d", resp.StatusCode)
	log.Printf("响应头: %+v", resp.Header)

	// 检查响应状态
	if resp.StatusCode != http.StatusOK {
		body, _ := ioutil.ReadAll(resp.Body)
		log.Printf("API错误响应: %s", string(body))
		return nil, fmt.Errorf("API error: status=%d, body=%s", resp.StatusCode, body)
	}

	// 读取响应体
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Printf("读取响应体失败: %v", err)
		return nil, err
	}

	log.Printf("原始响应体: %s", string(body))

	// 解析JSON响应
	var result BusinessResponse
	if err := json.Unmarshal(body, &result); err != nil {
		log.Printf("JSON解析失败: %v", err)
		return nil, fmt.Errorf("failed to parse API response: %v, body: %s", err, string(body))
	}

	log.Printf("解析后的响应数据:")
	for key, value := range result {
		log.Printf("  %s: %v", key, value)
	}
	log.Printf("=== API调用完成 ===\n")

	return result, nil
}

func generateResponse(template string, data BusinessResponse) string {
	result := template

	processedData := make(map[string]interface{})
	for key, value := range data {
		processedData[key] = value
	}

	// 处理用户状态描述
	if userState, ok := data["user_state"]; ok {
		if stateFloat, err := parseFloat(userState); err == nil {
			stateInt := int(stateFloat)
			switch stateInt {
			case 0:
				processedData["user_state_desc"] = "停机"
			case 1:
				processedData["user_state_desc"] = "正常"
			case 2:
				processedData["user_state_desc"] = "临时停机"
			default:
				processedData["user_state_desc"] = fmt.Sprintf("未知状态(%d)", stateInt)
			}
		}
	}

	// 处理使用时长（分钟转小时）
	if useTime, ok := data["use_time"]; ok {
		if timeFloat, err := parseFloat(useTime); err == nil {
			hours := timeFloat / 60.0
			processedData["use_time_hours"] = fmt.Sprintf("%.2f", hours)
		}
	}

	// 处理使用流量（MB转GB）
	if useFlow, ok := data["use_flow"]; ok {
		if flowFloat, err := parseFloat(useFlow); err == nil {
			gb := flowFloat / 1024.0
			processedData["use_flow_gb"] = fmt.Sprintf("%.2f", gb)
		}
	}

	// 替换模板中的占位符
	for key, value := range processedData {
		placeholder := "{{." + key + "}}"
		strValue := fmt.Sprintf("%v", value)
		result = strings.ReplaceAll(result, placeholder, strValue)
	}

	return result
}

func parseFloat(value interface{}) (float64, error) {
	switch v := value.(type) {
	case float64:
		return v, nil
	case float32:
		return float64(v), nil
	case int:
		return float64(v), nil
	case int64:
		return float64(v), nil
	case string:
		return parseFloatFromString(v)
	default:
		return 0, fmt.Errorf("unsupported type: %T", value)
	}
}

func parseFloatFromString(s string) (float64, error) {
	s = strings.TrimSpace(s)
	if s == "" {
		return 0, nil
	}

	if val, err := parseFloatValue(s); err == nil {
		return val, nil
	}

	return 0, fmt.Errorf("cannot parse %s as float", s)
}

func parseFloatValue(s string) (float64, error) {
	var result float64
	_, err := fmt.Sscanf(s, "%f", &result)
	return result, err
}

// 新增的验证相关函数
func loadAuthConfig() ClientAuthConfig {
	// 检查环境变量是否禁用验证
	if os.Getenv("MCP_AUTH_DISABLED") == "true" {
		log.Println("通过环境变量禁用客户端验证")
		return ClientAuthConfig{
			Enabled: false,
		}
	}

	// 尝试从配置文件加载，如果不存在则使用默认配置
	configFile := "auth_config.json"
	var config ClientAuthConfig

	if data, err := ioutil.ReadFile(configFile); err == nil {
		if json.Unmarshal(data, &config) == nil {
			log.Printf("已加载验证配置: %s", configFile)
		} else {
			log.Printf("配置文件格式错误，使用默认配置")
			config = getDefaultAuthConfig()
		}
	} else {
		log.Printf("未找到配置文件 %s，使用默认配置", configFile)
		config = getDefaultAuthConfig()
	}

	// 环境变量覆盖配置
	applyEnvironmentOverrides(&config)

	// 验证配置有效性
	validateAuthConfig(&config)

	return config
}

// 获取默认验证配置
func getDefaultAuthConfig() ClientAuthConfig {
	return ClientAuthConfig{
		Enabled:            true,
		RequireAPIKey:      true,
		RequireSignature:   false,
		RequireIPWhitelist: false,
		APIKeys: []string{
			"mcp_server_key_2024",
			"dev_test_key_123",
		},
		SignatureSecret: "mcp_server_secret_2024",
		IPWhitelist: []string{
			"127.0.0.1",
			"::1",
			"localhost",
		},
		MaxRetries:  3,
		BanDuration: 30, // 30分钟
	}
}

// 应用环境变量覆盖
func applyEnvironmentOverrides(config *ClientAuthConfig) {
	if apiKey := os.Getenv("MCP_API_KEY"); apiKey != "" {
		config.APIKeys = append(config.APIKeys, apiKey)
		log.Printf("从环境变量添加API密钥")
	}

	if secret := os.Getenv("MCP_SIGNATURE_SECRET"); secret != "" {
		config.SignatureSecret = secret
		log.Printf("从环境变量设置签名密钥")
	}

	if whitelistIPs := os.Getenv("MCP_IP_WHITELIST"); whitelistIPs != "" {
		ips := strings.Split(whitelistIPs, ",")
		for i, ip := range ips {
			ips[i] = strings.TrimSpace(ip)
		}
		config.IPWhitelist = append(config.IPWhitelist, ips...)
		log.Printf("从环境变量添加IP白名单: %v", ips)
	}
}

// 验证配置有效性
func validateAuthConfig(config *ClientAuthConfig) {
	if !config.Enabled {
		log.Println("客户端验证已禁用")
		return
	}

	if config.RequireAPIKey && len(config.APIKeys) == 0 {
		log.Println("警告：启用了API Key验证但未配置任何API密钥")
	}

	if config.RequireSignature && config.SignatureSecret == "" {
		log.Println("警告：启用了签名验证但未配置签名密钥")
	}

	if config.RequireIPWhitelist && len(config.IPWhitelist) == 0 {
		log.Println("警告：启用了IP白名单但未配置任何IP地址")
	}

	if config.MaxRetries <= 0 {
		config.MaxRetries = 3
		log.Println("修正：最大重试次数设置为默认值3")
	}

	if config.BanDuration <= 0 {
		config.BanDuration = 30
		log.Println("修正：封禁时长设置为默认值30分钟")
	}

	log.Printf("验证配置验证完成 - API Key: %t, 签名: %t, IP白名单: %t",
		config.RequireAPIKey, config.RequireSignature, config.RequireIPWhitelist)
}

// 创建客户端会话
func (am *AuthManager) createSession(clientIP string) *ClientSession {
	am.mutex.Lock()
	defer am.mutex.Unlock()

	sessionID := generateSessionID()
	session := &ClientSession{
		ID:              sessionID,
		ClientIP:        clientIP,
		AuthTime:        time.Now(),
		LastActivity:    time.Now(),
		RequestCount:    0,
		IsAuthenticated: false,
		FailedAttempts:  0,
		IsBanned:        false,
	}

	am.sessions[sessionID] = session
	return session
}

// 生成会话ID
func generateSessionID() string {
	timestamp := time.Now().UnixNano()
	data := fmt.Sprintf("session_%d_%d", timestamp, time.Now().UnixMilli())
	hash := sha256.Sum256([]byte(data))
	return hex.EncodeToString(hash[:])[:16]
}

// 验证客户端
func (am *AuthManager) authenticateClient(authInfo *AuthenticationInfo, clientIP string) (*ClientSession, error) {
	if !am.config.Enabled {
		// 验证未启用，创建默认会话
		session := am.createSession(clientIP)
		session.IsAuthenticated = true
		session.AuthMethod = "disabled"
		log.Printf("客户端验证已禁用，允许连接: %s", clientIP)
		return session, nil
	}

	// 检查IP是否被封禁
	if am.isIPBanned(clientIP) {
		return nil, fmt.Errorf("客户端IP已被封禁: %s", clientIP)
	}

	session := am.createSession(clientIP)

	// IP白名单验证
	if am.config.RequireIPWhitelist {
		if !am.isIPWhitelisted(clientIP) {
			am.recordFailedAttempt(session)
			return nil, fmt.Errorf("客户端IP不在白名单中: %s", clientIP)
		}
	}

	// API Key验证
	if am.config.RequireAPIKey {
		apiKey := ""

		// 优先从authInfo中获取API Key
		if authInfo != nil && authInfo.APIKey != "" {
			apiKey = authInfo.APIKey
			log.Printf("从初始化参数获取API Key")
		} else if envAPIKey := os.Getenv("MCP_API_KEY"); envAPIKey != "" {
			// 如果初始化参数中没有，则从环境变量获取
			apiKey = envAPIKey
			log.Printf("从环境变量获取API Key")
		}

		if apiKey == "" {
			am.recordFailedAttempt(session)
			return nil, fmt.Errorf("缺少API Key")
		}

		if !am.isValidAPIKey(apiKey) {
			am.recordFailedAttempt(session)
			return nil, fmt.Errorf("无效的API Key")
		}
		session.AuthMethod = "api_key"
	}

	// 签名验证
	if am.config.RequireSignature {
		if authInfo == nil || authInfo.Signature == "" || authInfo.Timestamp == "" {
			am.recordFailedAttempt(session)
			return nil, fmt.Errorf("缺少签名或时间戳")
		}

		if !am.verifySignature(authInfo) {
			am.recordFailedAttempt(session)
			return nil, fmt.Errorf("签名验证失败")
		}
		session.AuthMethod = "signature"
	}

	// 验证成功
	session.IsAuthenticated = true
	log.Printf("客户端验证成功: %s [%s]", clientIP, session.AuthMethod)
	return session, nil
}

// 检查IP是否在白名单中
func (am *AuthManager) isIPWhitelisted(ip string) bool {
	for _, whiteIP := range am.config.IPWhitelist {
		if ip == whiteIP || strings.Contains(ip, whiteIP) {
			return true
		}
	}
	return false
}

// 检查API Key是否有效
func (am *AuthManager) isValidAPIKey(apiKey string) bool {
	for _, validKey := range am.config.APIKeys {
		if apiKey == validKey {
			return true
		}
	}
	return false
}

// 验证签名
func (am *AuthManager) verifySignature(authInfo *AuthenticationInfo) bool {
	// 构建签名数据: timestamp + secret
	signData := authInfo.Timestamp + am.config.SignatureSecret
	hash := sha256.Sum256([]byte(signData))
	expectedSignature := hex.EncodeToString(hash[:])

	return authInfo.Signature == expectedSignature
}

// 检查IP是否被封禁
func (am *AuthManager) isIPBanned(ip string) bool {
	am.mutex.RLock()
	defer am.mutex.RUnlock()

	for _, session := range am.sessions {
		if session.ClientIP == ip && session.IsBanned {
			if time.Now().Before(session.BanUntil) {
				return true
			} else {
				// 封禁已过期，清除封禁状态
				session.IsBanned = false
				session.FailedAttempts = 0
			}
		}
	}
	return false
}

// 记录失败尝试
func (am *AuthManager) recordFailedAttempt(session *ClientSession) {
	am.mutex.Lock()
	defer am.mutex.Unlock()

	session.FailedAttempts++
	if session.FailedAttempts >= am.config.MaxRetries {
		session.IsBanned = true
		session.BanUntil = time.Now().Add(time.Duration(am.config.BanDuration) * time.Minute)
		log.Printf("客户端已被封禁: %s, 封禁至: %s", session.ClientIP, session.BanUntil.Format("2006-01-02 15:04:05"))
	}
}

// 获取客户端IP地址
func getClientIP() string {
	// 在MCP协议中，通常通过环境变量或其他方式获取客户端信息
	// 这里使用一个简化的实现
	if ip := os.Getenv("MCP_CLIENT_IP"); ip != "" {
		return ip
	}

	// 如果是通过网络连接，可以从连接中获取
	// 这里返回本地地址作为默认值
	return "127.0.0.1"
}

// 更新会话活动时间
func (am *AuthManager) updateSessionActivity(sessionID string) {
	am.mutex.Lock()
	defer am.mutex.Unlock()

	if session, exists := am.sessions[sessionID]; exists {
		session.LastActivity = time.Now()
		session.RequestCount++
	}
}

// 清理过期会话
func (am *AuthManager) cleanupExpiredSessions() {
	am.mutex.Lock()
	defer am.mutex.Unlock()

	expireTime := time.Now().Add(-24 * time.Hour) // 24小时过期
	for sessionID, session := range am.sessions {
		if session.LastActivity.Before(expireTime) {
			delete(am.sessions, sessionID)
			log.Printf("清理过期会话: %s", sessionID)
		}
	}
}
